# encoding=utf-8
import logging
import os
import threading
import traceback
import cv2
import numpy
import torch
from torch import nn
from Common.AppConfig import AppConfig
from torchvision import models
from ThirdPart.Grading.utils import circle_crop  # drop
from ThirdPart.Grading.demo import grad_dr_wrapper, grad_dr_visual_wrapper, grad_dr_wrapper_opti, grad_dr_visual_wrapper_opti
from torch.nn import functional as F
from torchvision.transforms import transforms


def get_kernel(kernel_len=16, nsig=10):  # nsig 标准差 ，kernlen=16核尺寸
    GaussianKernel = cv2.getGaussianKernel(kernel_len, nsig) \
                     * cv2.getGaussianKernel(kernel_len, nsig).T
    return GaussianKernel


class Gaussian_kernel(nn.Module):
    def __init__(self,
                 # device,
                 kernel_len, nsig=20):
        super(Gaussian_kernel, self).__init__()
        self.kernel_len = kernel_len
        kernel = get_kernel(kernel_len=kernel_len, nsig=nsig)  # 获得高斯卷积核
        kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0)  # 扩展两个维度
        # self.weight = nn.Parameter(data=kernel, requires_grad=False).to(device)
        self.weight = nn.Parameter(data=kernel, requires_grad=False)

        self.padding = torch.nn.ReplicationPad2d(int(self.kernel_len / 2))

    def forward(self, x):  # x1是用来计算attention的，x2是用来计算的Cs
        x = self.padding(x)
        # 对三个channel分别做卷积
        res = []
        for i in range(x.shape[1]):
            res.append(F.conv2d(x[:, i, :, :], self.weight))
        x_output = torch.cat(res, dim=0)
        return x_output


class MBGradeWrapperOpti:
    _inst = None
    _lock = threading.Lock()

    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = models.densenet121(pretrained=True)
        fc_feature = self.model.classifier.in_features
        self.model.classifier = nn.Linear(fc_feature, 5)
        self.model.load_state_dict(torch.load(AppConfig.RES_GRAD))
        self.model = self.model.to(self.device)
        self.model.eval()
        self.result_name = "cam-result"

        self.gaussian_kernel = Gaussian_kernel(11, 30).cuda()
        self.transform_kernel = transforms.Compose([
            transforms.ToTensor()
        ])

    def crop_image_from_mask(self, image_src: str):
        # load
        image = cv2.imread(image_src)

        # binary
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, binary_image = cv2.threshold(gray_image, 7, 255, cv2.THRESH_BINARY)

        # cal roi
        x, y, w, h = cv2.boundingRect(binary_image)
        center = (w // 2), (h // 2)
        radius = min(center)
        y = y + center[1] - radius
        x = x + center[0] - radius
        copy_image = image[y: y + 2 * radius, x: x + 2 * radius]

        # gen mask
        mask = numpy.zeros_like(copy_image)
        cv2.circle(mask, (radius, radius), radius, (1, 1, 1), -1)

        # exposure
        return copy_image * mask

    def crop_image_with_gaussian(self, data_in: numpy.ndarray):
        ori_image = cv2.resize(data_in, (224, 224)).astype(numpy.float32)
        with torch.no_grad():
            image_cuda = self.transform_kernel(ori_image).unsqueeze(0).cuda()
            out = numpy.transpose(self.gaussian_kernel(image_cuda).cpu().numpy(), (1, 2, 0))

        if out.shape != (224, 224, 3):
            out = out[0: 224, 0: 224]
        exposure = cv2.addWeighted(ori_image, 4, out, -4, 128)
        exposure = numpy.clip(exposure, 0, 255).astype(numpy.uint8)
        exposure = cv2.cvtColor(exposure, cv2.COLOR_BGR2RGB)
        return exposure

    def circle_crop(self, image_src: str):
        crop_mask = self.crop_image_from_mask(image_src)
        return self.crop_image_with_gaussian(crop_mask)

    def _preprocess(self, data_in: str):
        img_arr = self.circle_crop(data_in) / 255.0
        x = torch.from_numpy(img_arr[:, :, ::-1].astype(numpy.float32).transpose((2, 0, 1))).unsqueeze(0)
        return x, img_arr

    def _postprocess(self, cache_path: str, image_path: str, pred_arr):
        try:
            # grad cache
            grad_path = os.path.join(cache_path, "grade")
            os.makedirs(grad_path, exist_ok=True)

            grad_dr_visual_wrapper_opti(pred_arr, image_path, grad_path)
            return True, ""
        except Exception as e:
            except_str = traceback.format_exc()
            logging.exception(except_str)
            return False, except_str

    # @torch.no_grad()
    def infer(self, cache_path: str, image_path: str):
        trans_dim, img_arr = self._preprocess(image_path)
        trans_cuda = trans_dim.to(self.device)

        pred_arr = grad_dr_wrapper_opti(self.model, trans_cuda, self.device, img_arr)
        return self._postprocess(cache_path, image_path, pred_arr)

    def __new__(cls, *args, **kwargs):
        MBGradeWrapperOpti._lock.acquire()
        if cls._inst is None:
            cls._inst = super(MBGradeWrapperOpti, cls).__new__(cls)
        MBGradeWrapperOpti._lock.release()
        return cls._inst


class MBGradeWrapper:
    _inst = None
    _lock = threading.Lock()

    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = models.densenet121(pretrained=True)
        fc_feature = self.model.classifier.in_features
        self.model.classifier = nn.Linear(fc_feature, 5)
        self.model.load_state_dict(torch.load(AppConfig.RES_GRAD))
        self.model = self.model.to(self.device)
        self.model.eval()
        self.result_name = "cam-result"

    def _preprocess(self, data_in: str):
        img_arr = cv2.resize(circle_crop(data_in), (224, 224), interpolation=cv2.INTER_LINEAR) / 255.0
        x = torch.from_numpy(img_arr[:, :, ::-1].astype(numpy.float32).transpose((2, 0, 1))).unsqueeze(0)
        return x, img_arr

    def _postprocess(self, cache_path: str, image_path: str, pred_arr):
        try:
            file_name = os.path.splitext(os.path.basename(image_path))[0]

            # csv
            csv_path = os.path.join(cache_path, "csv")
            os.makedirs(csv_path, exist_ok=True)
            csv_file = os.path.join(csv_path, file_name + ".csv")

            # grad cache
            grad_path = os.path.join(cache_path, "grade")
            os.makedirs(grad_path, exist_ok=True)

            grad_dr_visual_wrapper(pred_arr, image_path, grad_path, csv_file)
            return True, ""
        except Exception as e:
            except_str = traceback.format_exc()
            logging.exception(except_str)
            return False, except_str

    # @torch.no_grad()
    def infer(self, cache_path: str, image_path: str):
        trans_dim, img_arr = self._preprocess(image_path)
        trans_cuda = trans_dim.to(self.device)

        pred_arr = grad_dr_wrapper(self.model, trans_cuda, self.device, img_arr)
        return self._postprocess(cache_path, image_path, pred_arr)

    def __new__(cls, *args, **kwargs):
        MBGradeWrapper._lock.acquire()
        if cls._inst is None:
            cls._inst = super(MBGradeWrapper, cls).__new__(cls)
        MBGradeWrapper._lock.release()
        return cls._inst
